import os
import shutil

def copy_files(source_folder, destination_folder):
    # 构造子文件夹路径
    images_test_path = os.path.join(destination_folder, 'images', 'test')
    images_train_path = os.path.join(destination_folder, 'images', 'train')
    images_val_path = os.path.join(destination_folder, 'images', 'val')
    labels_train_path = os.path.join(destination_folder, 'labels', 'train')
    labels_val_path = os.path.join(destination_folder, 'labels', 'val')

    # 检查并创建子文件夹
    for path in [images_test_path, images_train_path, images_val_path, labels_train_path, labels_val_path]:
        if not os.path.exists(path):
            os.makedirs(path)
            print(f"Created folder: {path}")
        else:
            print(f"Folder already exists: {path}")

    # 获取所有图片和标签文件
    images = []
    labels = []

    for file in os.listdir(source_folder):
        if file.endswith('.jpg') or file.endswith('.png'):
            images.append(file)
        elif file.endswith('.txt'):
            labels.append(file)

    # 复制所有图片到目标文件夹
    for image in images:
        # 直接将所有图片复制到 train、val 和 test 文件夹
        shutil.copy2(os.path.join(source_folder, image), os.path.join(images_train_path, image))
        print(f"Copied image: {image} to {images_train_path}")
        shutil.copy2(os.path.join(source_folder, image), os.path.join(images_val_path, image))
        print(f"Copied image: {image} to {images_val_path}")
        shutil.copy2(os.path.join(source_folder, image), os.path.join(images_test_path, image))
        print(f"Copied image: {image} to {images_test_path}")

    # 复制标签文件到相应的子文件夹
    for label in labels:
        # 假设所有标签文件都需要复制到 labels/train 和 labels/val
        shutil.copy2(os.path.join(source_folder, label), os.path.join(labels_train_path, label))
        print(f"Copied label: {label} to {labels_train_path}")
        shutil.copy2(os.path.join(source_folder, label), os.path.join(labels_val_path, label))
        print(f"Copied label: {label} to {labels_val_path}")

    print("文件拷贝成功")